from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
import logging
import gzip
import os
from postgresql.common.pg_exception import ErrCodeException
from common.const import ExecuteResultEnum
from pg_restore_service import PostgreRestoreService
from postgresql.common.util.pg_common_utils import PostgreCommonUtils
from postgresql.common.error_code import ErrorCode

LOGGER = logging.getLogger(__name__)


####################################### 解密模块 ##################################################
class Decryptor:
    def __init__(self, key):
        self._key = key.encode("utf-8")[:32]  # AES-256 requires a 32-byte key

    def process(self, source, target):
        """解密文件"""
        try:
            with open(source, "rb") as src:
                iv = src.read(16)  # 读取初始化向量
                cipher = AES.new(self._key, AES.MODE_CBC, iv)
                with open(target, "wb") as dst:
                    while chunk := src.read(8192):
                        decrypted_chunk = cipher.decrypt(chunk)
                        if src.tell() == os.path.getsize(source):  # 最后一块
                            decrypted_chunk = unpad(decrypted_chunk, AES.block_size)
                        dst.write(decrypted_chunk)
            LOGGER.info(f"Decrypted {source} -> {target}")
            return True
        except Exception as e:
            LOGGER.error(f"Failed to decrypt {source}: {str(e)}")
            return False
        
####################################### 解压缩模块 ##################################################
class Decompressor:
    def process(self, source, target):
        """解压缩文件"""
        try:
            with gzip.open(source, "rb") as src, open(target, "wb") as dst:
                dst.write(src.read())
            LOGGER.info(f"Decompressed {source} -> {target}")
            return True
        except Exception as e:
            LOGGER.error(f"Failed to decompress {source}: {str(e)}")
            return False
        
####################################### 恢复目录结构模块 ##################################################
class RestoreDirectory:
    def process(self, source, target):
        """恢复目录结构"""
        try:
            os.makedirs(os.path.dirname(target), exist_ok=True)
            with open(source, "rb") as src, open(target, "wb") as dst:
                dst.write(src.read())
            LOGGER.info(f"Restored {source} -> {target}")
            return True
        except Exception as e:
            LOGGER.error(f"Failed to restore {source}: {str(e)}")
            return False
        

####################################### 恢复数据主函数 ##################################################  
class RestoreManager:
    def __init__(self, encryption_key):
        self._decryptor = Decryptor(encryption_key)
        self._decompressor = Decompressor()
        self._restore_directory = RestoreDirectory()

    def _restore_file(self, source, target):
        """恢复单个文件（解密、解压缩、恢复目录结构）"""
        # 解密
        decrypted_file = source + ".dec"
        if not self._decryptor.process(source, decrypted_file):
            return False

        # 解压缩
        decompressed_file = decrypted_file + ".decompressed"
        if not self._decompressor.process(decrypted_file, decompressed_file):
            return False

        # 恢复目录结构
        if not self._restore_directory.process(decompressed_file, target):
            return False

        # 清理临时文件
        os.remove(decrypted_file)
        os.remove(decompressed_file)
        return True

    def restore_data(self, src_path, tgt_path, cache_path,wildcard, job_id=""):
        """恢复数据（逆流程）"""
        if not os.path.exists(src_path):
            LOGGER.error(f"Source directory does not exist: {src_path}")
            return False

        # 递归恢复目录
        for root, dirs, files in os.walk(src_path):
            relative_path = os.path.relpath(root, src_path)
            target_dir = os.path.join(tgt_path, relative_path)
            os.makedirs(target_dir, exist_ok=True)

            for file in files:
                if not file.endswith(".gz.enc"):
                    continue  # 跳过非备份文件
                source_file = os.path.join(root, file)
                target_file = os.path.join(target_dir, file[:-8])  # 去掉 .gz.enc 后缀
                if not self._restore_file(source_file, target_file):
                    return False

        # 恢复表空间
        restore_table_space_res = PostgreRestoreService.restore_table_space(src_path, cache_path, job_id)
        if restore_table_space_res != ExecuteResultEnum.SUCCESS:
            LOGGER.error("Restore table space failed!")
            raise ErrCodeException(restore_table_space_res, message="Restore table space failed")
        try:
            PostgreCommonUtils.copy_directory(src_path, tgt_path, wildcard=wildcard, job_id=job_id)
        except Exception as ex:
            LOGGER.error(f'Restore data failed, wildcard: "{wildcard}", source: {src_path}, target: {tgt_path}.')
            raise ErrCodeException(ErrorCode.RESTORE_DATA_FAILED, message="Restore data failed") from ex
        LOGGER.info(f"Restore completed successfully, job id: {job_id}.")
        return True